package com.bang.study.day05;

import java.util.TreeMap;

/**
 * @Auther: Bang
 * @Date: 2019/9/17 0017 08:54
 * @Description: 字典树 前缀树 基于TreeMap
 */
public class Trie {

    private class Node{

        /**
         * 是否为单词
         */
        public boolean isWord;

        /**
         * 存储下一个节点的
         */
        public TreeMap<Character, Node> next;

        public Node() {
            this.isWord = false;
            this.next = new TreeMap<>();
        }

        public Node(boolean isWord) {
            this();
            this.isWord = isWord;
        }
    }


    private Node root;

    private int size;


    public Trie() {
        this.root = new Node();
    }


    /**
     * 获取trie的大小
     * @return
     */
    public int getSize() {
        return size;
    }

    /**
     * 添加单词
     * @param word
     */
    public void add(String word) {
        Node curr = root;
        for (int i = 0; i < word.length(); i++) {
            char c = word.charAt(i);
            if (curr.next.get(c)==null) {
                curr.next.put(c, new Node());
            }
            curr = curr.next.get(c);
        }
        if (!curr.isWord) {
            curr.isWord = true;
            size++;
        }
    }

    /**
     * 是否包含某个单词
     * @param word
     * @return
     */
    public boolean contains(String word){
        Node curr = root;
        for (int i = 0; i <word.length() ; i++) {
            char c = word.charAt(i);
            if (curr.next.get(c)==null) {
                return false;
            }
            curr = curr.next.get(c);
        }
        return curr.isWord;
    }

    /**
     * 是否包含某个前缀
     * @param prefix
     * @return
     */
    public boolean containsPrefix(String prefix) {
        Node curr = root;
        for (int i = 0; i <prefix.length() ; i++) {
            char c = prefix.charAt(i);
            if (curr.next.get(c)==null) {
                return false;
            }
            curr = curr.next.get(c);
        }
        return true;
    }

    /**
     * 从trie中删除word
     * 1，如果单词是另一个单词的前缀，只需要把该word的最后一个节点的isWord的改成false
     * 2，如果单词的所有字母的都没有多个分支，删除整个单词
     * 3，如果单词的除了最后一个字母，其他的字母有多个分支，
     * @param word
     */
    public boolean remove(String word) {
        // 标识有多个分支的节点
        Node multiChildNode = null;
        // 标识有多个分支的索引
        int multiChildNodeIndex = -1;

        Node curr = root;
        for (int i = 0; i <word.length() ; i++) {
            char c = word.charAt(i);
            Node child = curr.next.get(c);
            if (child == null) {
                return false;
            }
            if (child.next.size() > 0) {
                multiChildNode = child;
                multiChildNodeIndex = i;
            }
            curr = child;
        }

        if (curr.next.size() > 0) {
            // 当前向下仍有节点
            if (curr.isWord){
                curr.isWord=false;
                size--;
                return true;
            }
            return false;
        }else{
            // 当前为最后一个节点
            if (multiChildNodeIndex == -1) {
                // 表示当前为单个单词
                root.next.remove(word.charAt(0));
                size--;
                return true;
            }

            //如果单词的除了最后一个字母，其他的字母有分支
            if (multiChildNodeIndex != word.length() - 1) {
                multiChildNode.next.remove(word.charAt(multiChildNodeIndex + 1));
                size--;
                return true;
            }
            return false;

        }


    }

    public static void main(String[] args) {
        Trie trie = new Trie();
        trie.add("hello");
        trie.add("hallo");
        trie.add("hcllo");
        trie.add("hdllo");

        System.out.println(trie.contains("hello"));
        System.out.println(trie.contains("hello1"));
        System.out.println(trie.containsPrefix("he"));
        System.out.println(trie.containsPrefix("ae"));
        System.out.println(trie.remove("hello"));
        System.out.println(trie.contains("hello"));


    }


}
